import torch
from torch.utils.tensorboard import SummaryWriter


model = torch.nn.LSTM(10, 20, 2)  # 参数：输入维度，隐藏层维度，层数
input = torch.randn(5, 3, 10)  # 参数：序列长度，batch，输入维度
h0 = torch.randn(2, 3, 20)  # h0 代表初始隐藏层，参数：层数，batch，隐藏层维度
c0 = torch.randn(2, 3, 20)  # c0 代表初始记忆单元
output, (hn, cn) = model(input, (h0, c0))
print(output.size())

with SummaryWriter(log_dir='') as sw:  # 实例化 SummaryWriter ,可以自定义数据输出路径
    sw.add_graph(model, input)  # 输出网络结构图
    sw.close()  # 关闭  sw